# 模型文件地址
import os
from . import path_utils


def _dataset_info_(dataset_name, n_cls=None):
    info = {'dataset_dir': os.path.join(path_utils.dataset_base_dir, dataset_name),
            'n_cls': n_cls}
    return info


def _model_file_path_(model_name):
    model_dir = os.path.join(path_utils.model_base_dir, model_name)
    model_file_name = "pytorch_model.bin"
    config_file_name = "config.json"
    vocab_file_name = "vocab.txt"
    info = {'model_dir': model_dir, 'model_path': None, 'config_path': None, 'vocab_path': None}
    if os.path.exists(os.path.join(model_dir, model_file_name)):
        info['model_path'] = os.path.join(model_dir, model_file_name)
    if os.path.exists(os.path.join(model_dir, config_file_name)):
        info['config_path'] = os.path.join(model_dir, config_file_name)
    if os.path.exists(os.path.join(model_dir, vocab_file_name)):
        info['vocab_path'] = os.path.join(model_dir, vocab_file_name)
    return info


model_file_info_map = {
    'bert_base_chinese': _model_file_path_(model_name='bert_base_chinese'),
    'ernie': _model_file_path_(model_name='ernie'),
    'roberta_tiny_clue': _model_file_path_(model_name='roberta_tiny_clue'),
    'roberta_tiny_pair': _model_file_path_(model_name='roberta_tiny_pair'),
}

dataset_dir_map = {
    'sort': _dataset_info_('d101_sort'),
    'THUCNews_title': _dataset_info_('d102_THUCNews_title', n_cls=10),
    'ccks2018': _dataset_info_('d104_ccks2018_sim', n_cls=2),
    'lcqmc': _dataset_info_('d105_lcqmc_sim', n_cls=2),
    'snli': _dataset_info_('d106_snli_sim', n_cls=3),
    'ocnli': _dataset_info_('d107_ocnli_sim', n_cls=3),
}


